import json
import re
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM


# ---------------------- 配置参数 ----------------------
MODEL_PATH = "/data2/modelscope/hub/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"  # 本地模型路径
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
INPUT_DATA_PATH = "train_cf/src/20250623/pt_train_cf.json"  # 输入文章数据集路径（JSON格式）
# OUTPUT_DATA_PATH = "train_cf/src/20250623/output—800.json"  # 输出SFT数据集路径
OUTPUT_VALID_PATH = "train_cf/src/20250623/output—800.json"  # 有效样本输出路径
OUTPUT_INVALID_PATH = "train_cf/src/20250623/invalud-output—800.json"  # 无效样本输出路径

# ---------------------- 加载本地模型与分词器 ----------------------
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        device_map="auto",
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32
    )
except Exception as e:
    raise RuntimeError(f"加载本地模型失败：{e}，请检查模型路径是否正确")

# ---------------------- 从JSON文件读取输入数据 ----------------------
def load_articles_from_json(file_path: str) -> list:
    """从JSON文件加载文章数据集"""
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, list):
                return data
            else:
                raise ValueError("JSON文件格式错误：数据应为列表格式（每篇文章为字符串）")
    except FileNotFoundError:
        raise FileNotFoundError(f"未找到输入文件：{file_path}，请检查路径是否正确")
    except json.JSONDecodeError:
        raise ValueError(f"JSON文件解析失败：{file_path}，请检查文件格式是否正确")

# ---------------------- 通用提示模板（扩展后） ----------------------
PROMPT_TEMPLATE = """
任务说明：
你需要根据输入的文章内容，提取其中“第一作者（我/作者）亲身经历的事件”，生成符合以下要求的对话式样本：
- 样本格式：用户提问（引导个人经历）→ 作者回答（包含具体事件细节）；
- 用户提问需围绕“第一作者的个人经历”设计（如“你能和我聊聊你做试管婴儿的经历吗？”“你在香港租房时遇到过哪些困难？”“在北京买房的过程中有什么特别的经历？”）；
- 作者回答必须包含：
  ① 第一人称行动（如“我调整了促排时间”“我在香港租了XX公寓”“我在北京看了XX小区的房子”）；
  ② 时间/地点/人物（如“2024年3月在深圳出差”“香港铜锣湾XX大厦”“北京朝阳区XX小区”）；
  ③ 情感反馈（如“我当时很焦虑”“租房时被中介骗了很生气”“买房后终于有了归属感”）；
  ④ 社会背景关联（如“高龄试管的医学挑战”“香港租房的高昂成本”“北京楼市的政策变化”）；
- 若输入文章中无“第一作者亲身经历”（仅描述他人故事或客观事实），则输出“无相关个人经历”。

输入数据（需替换为实际文章内容）：
{input_article}

生成规则：
1. 用户提问需自然引导作者分享个人经历（如“能详细说说你做试管这一年最难忘的事吗？”“在香港租房时你最担心的问题是什么？”）；
2. 作者回答必须基于输入数据中的具体事件，禁止虚构；
3. 若输入数据中无“我/作者”的具体行动、时间、地点等细节，则输出“无相关个人经历”。

示例引导（输入数据+期望输出）：
【输入数据示例】：“去年崔姐因妹夫车祸去世返家，我收留了她5岁的女儿小诺。小诺刚来时总躲在角落哭，我每天陪她玩积木、讲睡前故事。有天她突然说：‘阿姨，你家比我家暖和。’我才知道她以前住在没暖气的破房子里。”

【期望输出示例】：
用户提问：“你能和我聊聊你和崔姐最近一次关于她女儿的互助经历吗？”
作者回答：“去年冬天，崔姐因为妹夫出车祸去世，家里只剩她和5岁的小诺，房子还卖不出去，就带着孩子来北京投奔我。小诺刚到我家的时候特别胆小，总躲在沙发角落哭，饭也不肯吃。我每天晚上陪她玩积木，给她讲《小熊绘本》的故事，慢慢她才愿意跟我说话。有天半夜她突然抱着我的胳膊说：‘阿姨，你家比我家暖和多了。’我才知道她以前住的破房子连暖气都没有，冬天冷得直打颤。那一刻我特别心疼——她才5岁，就要承受这些。”
"""

# ---------------------- 生成函数 ----------------------
def generate_personal_experience(input_article: str) -> dict:
    """调用本地模型生成个人经历样本"""
    # 动态填充输入数据到提示模板
    prompt = PROMPT_TEMPLATE.format(input_article=input_article)
    
    # 编码输入
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
    
    # 模型推理（生成1个样本）
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=800,  # 控制生成长度
            num_return_sequences=1,
            temperature=0.7,  # 控制随机性（0-1，越低越确定）
            do_sample=True
        )
    
    # 解码生成结果
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # 提取“用户提问→作者回答”部分（根据提示模板的结构，生成结果通常在提示末尾）
    try:
        # 假设生成结果中“用户提问：”和“作者回答：”是固定标记（需根据实际提示调整）
        user_prompt_start = generated_text.find("用户提问：")
        user_prompt_end = generated_text.find("作者回答：", user_prompt_start)
        author_response_start = user_prompt_end + len("作者回答：")
        author_response_end = generated_text.rfind("。")  # 简单截断（可根据实际调整）
        
        user_question = generated_text[user_prompt_start+5:user_prompt_end].strip()
        author_answer = generated_text[author_response_start:author_response_end].strip()
        
        # 验证是否包含核心要素（简化版）
        has_first_person = any(keyword in author_answer for keyword in ["我", "我做了", "我调整了", "我打了", "我租了", "我买了"])
        has_detail = any(keyword in author_answer for keyword in ["时间", "地点", "人物", "情感", "社会背景"])
    
        if has_first_person and has_detail:
            return {"用户提问": user_question, "作者回答": author_answer}
        else:
            return {"用户提问": "", "作者回答": "无相关个人经历"}
    except:
        return {"用户提问": "", "作者回答": "无相关个人经历"}

    # # 优化提取逻辑：使用更灵活的正则表达式匹配（兼容换行符、标点）
    # user_prompt_pattern = r"用户提问：(.*?)作者回答："
    # author_response_pattern = r"作者回答：(.*?)\n\n"
    
    # # 提取用户提问（去除前缀“用户提问：”）
    # user_question_match = re.search(user_prompt_pattern, generated_text, re.DOTALL)
    # user_question = user_question_match.group(1).strip() if user_question_match else ""
    
    # # 提取作者回答（去除前缀“作者回答：”）
    # author_answer_match = re.search(author_response_pattern, generated_text, re.DOTALL)
    # author_answer = author_answer_match.group(1).strip() if author_answer_match else ""
    
    # # 判断是否为有效样本（包含第一人称行动、时间/地点/人物、情感、社会背景）
    # has_first_person = any(keyword in author_answer for keyword in ["我", "我做了", "我调整了", "我打了", "我租了", "我买了"])
    # has_detail = any(keyword in author_answer for keyword in ["时间", "地点", "人物", "情感", "社会背景"])
    
    # # 返回结果及有效性标记
    # return {
    #     "用户提问": user_question,
    #     "作者回答": author_answer,
    #     "is_valid": has_first_person and has_detail  # 有效性标记
    # }

# ---------------------- 批量生成并保存结果 ----------------------
def main():
    # 加载输入数据
    try:
        articles = load_articles_from_json(INPUT_DATA_PATH)
        print(f"成功加载{len(articles)}篇文章数据")
    except Exception as e:
        print(f"加载输入数据失败：{e}")
        return
    
    # 批量生成样本
    dataset = []
    for idx, article in enumerate(articles, 1):
        print(f"正在处理第{idx}/{len(articles)}篇文章...")
        print(f"文章为：{article}")
        sample = generate_personal_experience(article)
        print(f"生成的数据为{sample}")
        if sample["作者回答"] != "无相关个人经历":
            dataset.append(sample)
        else:
            print(f"第{idx}文章中无相关个人经历")
    
    # 保存结果（用于Llama-Factory训练）
    with open(OUTPUT_DATA_PATH, "w", encoding="utf-8") as f:
        json.dump(dataset, f, ensure_ascii=False, indent=2)
    
    print(f"生成完成！有效样本数：{len(dataset)}，保存路径：{OUTPUT_DATA_PATH}")

    # # 初始化分类列表
    # valid_samples = []
    # invalid_samples = []
    
    # # 批量生成样本并分类
    # for idx, article in enumerate(articles, 1):
    #     print(f"正在处理第{idx}/{len(articles)}篇文章...")
    #     sample = generate_personal_experience(article)
        
    #     # 根据有效性分类
    #     if sample["is_valid"]:
    #         valid_samples.append({"用户提问": sample["用户提问"], "作者回答": sample["作者回答"]})
    #     else:
    #         invalid_samples.append({
    #             "输入文章": article,
    #             "生成结果": {
    #                 "用户提问": sample["用户提问"],
    #                 "作者回答": sample["作者回答"]
    #             }
    #         })
    # # 保存有效样本
    # with open(OUTPUT_VALID_PATH, "w", encoding="utf-8") as f:
    #     json.dump(valid_samples, f, ensure_ascii=False, indent=2)
    # print(f"有效样本保存完成！数量：{len(valid_samples)}，路径：{OUTPUT_VALID_PATH}")
    
    # # 保存无效样本
    # with open(OUTPUT_INVALID_PATH, "w", encoding="utf-8") as f:
    #     json.dump(invalid_samples, f, ensure_ascii=False, indent=2)
    # print(f"无效样本保存完成！数量：{len(invalid_samples)}，路径：{OUTPUT_INVALID_PATH}")

if __name__ == "__main__":
    main()